#lang racket/base
(provide
 expand-to-core/shallow
 expand-def/shallow
 expand-expression
 )

(require
 (for-template
  "literals.rkt")
 racket/set
 racket/list
 syntax/apply-transformer
 syntax/parse
 (only-in syntax/parse [attribute @])
 "transformers/infix.rkt"
 "../util/stx.rkt")

;; id -> (or (values 'normal #f) (values 'macro procedure?))
(define (identifier-meaning id #:src [src id] #:ctx [ctx* '()])
  (define val (syntax-local-value id (λ () #f) ctx*))
  (cond
    [(not val)
     (values 'normal #f)]
    [(procedure? val)
     (values 'macro val)]
    [else
     (raise-syntax-error #f "unrecognized identifier" src)]))

;; stx -> stx
;| recursively expand the outer part of the syntax until
;| it becomes a core form, then return that core form.
(define (expand-to-core/shallow stx #:ctx [ctx* '()])

  (define (local-apply-here f)
    (local-apply-transformer f stx 'expression ctx*))

  (syntax-parse stx
    ; convert infix notation into prefix
    [(t ...)
     #:when (eq? (syntax-property stx 'paren-shape) #\{)
     (expand-to-core/shallow
      (rewrite-infix->prefix (@ t) #:src stx)
      #:ctx ctx*)]

    ; possibly expand variables
    [x:id
     (define-values [meaning proc]
       (identifier-meaning #'x #:ctx ctx*))
     (case meaning
       [(normal) #'x]
       [(macro)
        (expand-to-core/shallow
         (local-apply-here proc)
         #:ctx ctx*)])]

    ; stop at core forms
    [(hd:id . _)
     #:when (set-member? core-literal-ids/set #'hd)
     stx]

    ; expand (hd . tl) when hd is a macro
    [(hd:id . tl)
     #:do [(define-values [meaning proc]
             (identifier-meaning #'hd #:ctx ctx*))]
     #:when (eq? meaning 'macro)
     (expand-to-core/shallow
      (local-apply-here proc)
      #:ctx ctx*)]

    ; introduce #%sp:app otherwise
    [(hd . tl)
     (expand-to-core/shallow
      (syntax/loc stx (#%sp:app hd . tl))
      #:ctx ctx*)]

    ; rewrite datum into #%sp:datum
    [x
     #:when (valid-datum? (syntax-e #'x))
     (syntax/loc stx
       (#%sp:datum x))]

    [_ (raise-syntax-error #f "invalid syntax" stx)]))

;; stx -> #'(#%sp:def x e)
;| recursively expand the outer part of the syntax until it becomes
;| a core form. if this core form is not `#%sp:def`, assume it is an
;| expression and tranform it into a `#%sp:def` binding a random
;| temporary.
(define (expand-def/shallow stx #:ctx [ctx* '()])
  (syntax-parse (expand-to-core/shallow stx #:ctx ctx*)
    #:literal-sets [core-literals]
    [(hd:#%sp:def pat body) this-syntax]
    [(hd:#%sp:def pat body . stuff)
     (raise-syntax-error #f "too many terms in #%sp:def" stx)]
    [expr
     #:with [tmp] (generate-temporaries #'[_])
     #'(#%sp:def tmp expr)]))

;; stx -> stx
;| fully expand an expression. if it succeeds, the resulting syntax is
;| guarunteed to be well formed. keyword arguments are rewritten to occur
;| after non-keyword arguments.
(define (expand-expression stx #:ctx [ctx* '()])
  (define (>>expr x [ctxs ctx*]) (expand-expression x #:ctx ctxs))
  (syntax-parse (expand-to-core/shallow stx #:ctx ctx*)
    #:literal-sets [core-literals]
    [x:id #'x]
    [(hd:#%sp:datum ~! d)
     #:fail-unless (valid-datum? (syntax-e #'d)) "invalid data literal"
     (stx/loc (hd d))]
    [(hd:#%sp:primop ~! name:string)
     (stx/loc (hd name))]

    ; (#%sp:app fn [arg | #:kw karg] ...)
    [(hd:#%sp:app ~!
                  fn:expr
                  {~alt {~seq k:keyword karg:expr}
                        arg:expr}
                  ...)
     #:fail-when (check-duplicates (@ k) #:key syntax-e)
     "duplicate keyword in application"
     #:with fn* (>>expr #'fn)
     #:with [arg* ...] (map >>expr (@ arg))
     #:with [karg* ...] (map >>expr (@ karg))
     #:with [[k+karg* ...] ...] #'[[k karg*] ...]
     (stx/loc (hd fn* arg* ... k+karg* ... ...))]

    ; (#%sp:λ ([id | #:kw kid] ...) body)
    [(hd:#%sp:λ ~!
                ({~alt {~seq k:keyword karg:id}
                      arg:id}
                 ...)
                body:expr)
     #:fail-when (check-duplicates (@ k) #:key syntax-e)
     "duplicate keyword in function parameters"
     #:do [(define ctx (syntax-local-make-definition-context))
           (define (>>ctx s) (internal-definition-context-introduce ctx s))
           (syntax-local-bind-syntaxes (@ arg) #f ctx)
           (syntax-local-bind-syntaxes (@ karg) #f ctx)]
     #:with [arg*  ...] (map >>ctx (@ arg))
     #:with [karg* ...] (map >>ctx (@ karg))
     #:with body* (>>expr (>>ctx #'body) (cons ctx ctx*))
     #:with [[k+karg* ...] ...] #'[[k karg*] ...]
     (stx/loc (hd (arg* ... k+karg* ... ...) body*))]

    ; (#%sp:let (x rhs) ... body)
    [(hd:#%sp:let ~! (x:id rhs:expr) ... body:expr)
     #:do [(define ctx (syntax-local-make-definition-context))
           (define (>>ctx s) (internal-definition-context-introduce ctx s))
           (define (>>expr* s) (>>expr (>>ctx s) (cons ctx ctx*)))
           (syntax-local-bind-syntaxes (@ x) #f ctx)]
     #:with [x*   ...] (map >>ctx (@ x))
     #:with [rhs* ...] (map >>expr* (@ rhs))
     #:with body* (>>expr* #'body)
     (stx/loc (hd (x* rhs*) ... body*))]

    [u (raise-syntax-error #f "invalid expression" #'u)]))
